{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KP-ABE\n",
    "非常典型的KP-ABE构造方案，来自论文Attribute-Based Encryption for Fine-Grained Access Control of Encrypted Data，为了避免不准确的翻译，下面的方案构造过程中英文全是论文中的原话，中文是我自己的理解。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let $ G_1 $  be a bilinear group of prime order $ p $ ,and let $g$ be a generator of $ G_1 $ .In addition, let $ e:G_1\\times G_1\\rightarrow G_2 $ denote the bilinear map. A security parameter,$ k $,will determine the size of the groups.We also define the Lagrange coefficient $ \\Delta _{i,S} $ for $ i\\in \\mathbb{Z} _p $ and a set,$S$, of elements in $ \\mathbb{Z} _p:\\Delta _{i,s}\\left( x \\right) =\\prod_{j\\in S,j\\ne i}{\\frac{x-j}{i-j}} $ We will associate each attribute with a unique element in $ \\mathbb{Z} _{p}^{*} $.\n",
    "Our construction follows."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "从上面的原文中我们注意到$ \\Delta _{i,S}\\left( x \\right) =\\prod_{j\\in S,j\\ne i}{\\frac{x-j}{i-j}} $其实就是拉格朗日插值基函数，而且在下面秘密恢复的过程中需要用到，所以在下面的代码中给出它的实现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from charm.toolbox.pairinggroup import PairingGroup, G1, GT, ZR, pair as e\n",
    "from functools import reduce\n",
    "from node import *\n",
    "import numpy as np\n",
    "\n",
    "group = PairingGroup('SS512')\n",
    "g = group.random(G1)\n",
    "\n",
    "def lagrange_interpolate_base_fun(i, s, x):\n",
    "    \"\"\" 计算由 i 和 S (i 属于 X) 确定的拉格朗日插值基函数在x处的函数值\"\"\"\n",
    "    ans = group.init(ZR, 1)\n",
    "    for j in s:\n",
    "        if i != j:\n",
    "            ans *= (x - j) / (i - j)\n",
    "    return ans"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup\n",
    "Define the universe of attributes $ \\mathcal{U} =\\left\\{ 1,2,...,n \\right\\}  $. Now,for each attribute $ i\\in \\mathcal{U}  $, choose a number $t_i$ uniformly at random from $ \\mathbb{Z} _p $.Finally,choose $y$ uniformly at random in $ \\mathbb{Z} _p $.The published public parameters PK are\n",
    "$$\n",
    "T_1=g^{t_1},...,T_{\\left| \\mathcal{U} \\right|}=g^{t_{\\left| \\mathcal{U} \\right|}},Y=e\\left( g,g \\right) ^y\n",
    "$$\n",
    "The master key MK is:\n",
    "$$\n",
    "t_1,...,t_{\\left| \\mathcal{U} \\right|},y\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在实现过程中把$n$作为参数传入，并且属性改为$ \\mathcal{U} =\\left\\{ 0,1,...,n-1 \\right\\}  $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup(n):\n",
    "    t = [group.random(ZR) for _ in range(n)]\n",
    "    T = [g ** ti for ti in t]\n",
    "    y = group.random(ZR)\n",
    "    Y = e(g, g) ** y\n",
    "    return {'PK': {'T': T, 'Y': Y}, 'MK': {'t': t, 'y': y}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Encryption($M$, $\\gamma$, PK)\n",
    "To encrypt a message $M \\in G_2$ under a set of attributes $\\gamma$, choose a random value $ s\\in \\mathbb{Z} _p $ and publish the ciphertext as:\n",
    "$$\n",
    "E=\\left( \\gamma ,E'=MY^s,\\left\\{ E_i=T_{i}^{s} \\right\\} _{i\\in \\gamma} \\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "为了便于后续的实现，在encryption函数上多传入一个参数，该参数表示全局属性的大小，并且将上述的密文格式改为\n",
    "$$\n",
    "E=\\left( \\gamma ,E',MY^s,\\left\\{ E_{i\\in \\mathcal{U}}=\\begin{cases}\n",
    "\tT_{i}^{s}&\t\ti\\in \\gamma\\\\\n",
    "\tNone&\t\totherwise\\\\\n",
    "\\end{cases} \\right\\} \\right)\n",
    "$$\n",
    "显然这样的做法，并不会影响解密，但在实际过程中暴露了全局属性，但是为了方便测试，我们选择后者的实现方式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encryption(M, gamma, PK, U_size):\n",
    "    \"\"\" M is the element of G2 \"\"\"\n",
    "    s = group.random(ZR)\n",
    "    E = [None for _ in range(U_size)]\n",
    "    for i in gamma:\n",
    "        E[i] = PK['T'][i] ** s\n",
    "    return {'gamma': gamma, 'E_minute': M * (PK['Y'] ** s), 'E': E}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Key Generation($ \\mathcal{T}  $, MK)\n",
    "The algorithm outputs a key that enables the user to decrypt a message encrypted under a set of attributes $\\gamma$ if and only if $ \\mathcal{T} \\left( \\gamma \\right) =1 $. The algorithm proceeds as follows. First choose a polynomial $q_x$ for each node $x$ (including the leaves) in the tree $\\mathcal{T}$. These polynomials are chosen in the following way in a top-down manner, starting from the root node $r$. \\\n",
    "For each node $x$ in the tree, set the degree $d_x$ of the polynomial $q_x$ to be one less than the threshold value $k_x$ of that node, that is, $d_x = k_x - 1$. Now, for the root node $r$, set $q_r(0) = y$ and $d_r$ other points of the polynomial $q_r$ randomly to define it completely. For any other node $x$, set $ q_x\\left( 0 \\right) =q_{parent\\left( x \\right)}\\left( index\\left( x \\right) \\right)  $ and choose $d_x$ other points randomly to completely define $q_x$. \\\n",
    "Once the polynomials have been decided, for each leaf node $z$, we give the following secret value to the user:\n",
    "$$\n",
    "D_x=g^{\\frac{q_x\\left( 0 \\right)}{t_i}}\\,\\,\\mathrm{where} i=att\\left( x \\right)\n",
    "$$\n",
    "The set of above secret values is the decryption key $D$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def secret_share(root: Node, s):\n",
    "    \"\"\" 秘密共享 \"\"\"\n",
    "    root.secret = s\n",
    "    if not root.is_leaf():\n",
    "        coef = np.append([group.random(ZR) for _ in range(root.gate.k - 1)], s)\n",
    "        for x, c in enumerate(root.children, 1):\n",
    "            # 由于是在Zp上进行运算，所以把x初始化为Zp上的元素\n",
    "            secret_share(c, np.polyval(coef, group.init(ZR, x)))\n",
    "\n",
    "def key_generation(T, MK):\n",
    "    secret_share(T, MK['y'])\n",
    "    # 获取所有的叶子结点\n",
    "    leafs = []\n",
    "    def dfs(root: Node):\n",
    "        if root.is_leaf():\n",
    "            leafs.append(root)\n",
    "        else:\n",
    "            for c in root.children:\n",
    "                dfs(c)\n",
    "    dfs(T)\n",
    "    D = [None for _ in range(len(leafs))]\n",
    "    for leaf in leafs:\n",
    "        D[leaf.attr] = g ** (leaf.secret / MK['t'][leaf.attr])\n",
    "    return D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decryption($E$, $D$)\n",
    "We specify our decryption procedure as a recursive algorithm. For ease of exposition we present the simplest form of the decryption algorithm and discuss potential performance improvements in the next subsection.\n",
    "We first define a recursive algorithm DecryptNode($E$,$D$,$x$) that takes as input the ciphertext $ E=\\left( \\gamma ,E',\\left\\{ E_i \\right\\} _{i\\in \\gamma} \\right)  $, the private key $D$(we assume the access tree $\\mathcal{T}$ is embedded in the private key), and a node $x$ in the tree. It outputs a group element of $G_2$ or $ \\bot  $.\n",
    "Let $i=$ att($x$).If the node $x$ is a leaf node then:\n",
    "$$\n",
    "\\mathrm{DecryptNode}\\left( E,D,x \\right) =\\begin{cases}\n",
    "\te\\left( D_x,E_i \\right) =e\\left( g^{\\frac{q_x\\left( 0 \\right)}{t_i}},g^{s\\cdot t_i} \\right) =e\\left( g,g \\right) ^{s\\cdot q_x\\left( 0 \\right)}\\,\\,\\mathrm{if} i\\in \\gamma\\\\\n",
    "\t\\bot \\,\\,\\mathrm{otherwise}\\\\\n",
    "\\end{cases}\n",
    "$$\n",
    "We now consider the recursive case when $x$ is a non-leaf node. The algorithm DecryptNode($E$,$D$,$x$) then proceeds as follows: For all nodes $z$ that are children of $x$,it calls DecryptNode($E$,$D$,$z$) and stores the output as $F_z$.Let $S_x$ be an arbitrary $k_x$-sized set of child nodes $z$ such that $ \\mathrm{F}_z\\ne \\bot  $.If no such set exists then the node was not satisfied and the function returns $\\bot$.\n",
    "Otherwise,we compute:\n",
    "\\begin{align*}\n",
    "F_x&=\\prod_{z\\in S_x}{{F_z}^{\\Delta _{i,S_{x}^{'}}\\left( 0 \\right)}}, \\mathrm{where} \n",
    "\\begin{cases}\n",
    "\ti=index\\left( z \\right)\\\\\n",
    "\tS_{x}^{'}=\\left\\{ index\\left( z \\right) :z\\in S_x \\right\\}\\\\\n",
    "\\end{cases}\n",
    "\\\\\n",
    "&=\\prod_{z\\in S_x}{\\left( e\\left( g,g \\right) ^{s\\cdot q_z\\left( 0 \\right)} \\right) ^{\\Delta _{i,S_{x}^{'}}\\left( 0 \\right)}}\n",
    "\\\\\n",
    "&=\\prod_{z\\in S_x}{\\left( e\\left( g,g \\right) ^{s\\cdot q_{parent\\left( z \\right)}\\left( index\\left( z \\right) \\right)} \\right) ^{\\Delta _{i,S_{x}^{'}}\\left( 0 \\right)}}\\,\\,\\left( \\mathrm{by-} \\mathrm{construction} \\right) \n",
    "\\\\\n",
    "&=\\prod_{z\\in S_x}{e\\left( g,g \\right) ^{s\\cdot q_x\\left( i \\right) \\cdot \\Delta _{i,S_{x}^{'}}\\left( 0 \\right)}}\n",
    "\\\\\n",
    "&=e\\left( g,g \\right) ^{s\\cdot q_x\\left( 0 \\right)}\\,\\,\\left( \\mathrm{using-} \\mathrm{polynomial-} \\mathrm{interpolate} \\right)\n",
    "\\end{align*}\n",
    "and return the result.\n",
    "Now that we have defined our function DecryptNode,the decryption algorithm sim-ply calls the function on the root of the tree.We observe that DecryptNode($E$,$D$,$r$)=$ e\\left( g,g \\right) ^{ys} $=$Y^s$ if and only if the ciphertext satisfies the tree.Since,$ E'=MY^s $ the decryption algorithm simply divides out $ Y^s $ and recovers the message $M$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "解密算法其实就是直接实现DecryptNode即可，需要注意的是，传入的x节点以及所有的子节点的秘密值都为None，所以在调用前，可以先将所有的节点的值清空，这样做的目的是模拟真实情况，因为我们不可能拿着已分享秘密值的访问树来恢复秘密。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decrypt_node(E, D, x: Node):\n",
    "    if x.is_leaf():\n",
    "        if x.attr in E['gamma']:\n",
    "            x.secret = e(D[x.attr], E['E'][x.attr])\n",
    "            return True\n",
    "    else:\n",
    "        points = []\n",
    "        for index, z in enumerate(x.children, 1):\n",
    "            if decrypt_node(E, D, z):\n",
    "                # 将index初始化为Zp上的元素，方便后续插值恢复\n",
    "                points.append((group.init(ZR, index), z.secret))\n",
    "            if len(points) == x.gate.k:\n",
    "                X, Y = zip(*points)\n",
    "                # 计算乘积项 Fz ^ {DELTA_i_s(0)}\n",
    "                base_list = [Y[i] ** lagrange_interpolate_base_fun(X[i], X, group.init(ZR, 0)) for i in range(x.gate.k)]\n",
    "                # 将乘积因子连乘并赋值到根节点\n",
    "                x.secret = reduce(lambda e1, e2: e1 * e2, base_list)\n",
    "                return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "明文消息：\n",
      " [7769403840989686172657524729657194583930073095868442011218058703156222093195487595301716244479506586513109231127210198567678834119189360728199232812729111, 1088892259340777079722636119863550322423360294243847837640260608941191271949311910623910031538728607280141245191991992319003299908021496260927716691513666]\n",
      "----------------------------------- Access Tree Structure -----------------------------------\n",
      "(2, 3): None\n",
      "\t(3, 3): None\n",
      "\t\t0: None\n",
      "\t\t1: None\n",
      "\t\t2: None\n",
      "\t6: None\n",
      "\t(1, 3): None\n",
      "\t\t3: None\n",
      "\t\t4: None\n",
      "\t\t5: None\n",
      "\n",
      "---------------------------------------------------------------------------------------------\n",
      "解密成功！解密消息：\n",
      " [7769403840989686172657524729657194583930073095868442011218058703156222093195487595301716244479506586513109231127210198567678834119189360728199232812729111, 1088892259340777079722636119863550322423360294243847837640260608941191271949311910623910031538728607280141245191991992319003299908021496260927716691513666]\n"
     ]
    }
   ],
   "source": [
    "# 假设属性全集为 0 ... 7, 并且访问结构如下\n",
    "u_size = 7 # 全局属性大小\n",
    "gamma = [2, 0, 1, 3] # 文档属性\n",
    "M = group.random(GT) # 消息M\n",
    "print('明文消息：\\n', M)\n",
    "\n",
    "t33 = Node.threshold_node(Gate(3, 3), [Node.attr_node(0), Node.attr_node(1), Node.attr_node(2)])\n",
    "t13 = Node.threshold_node(Gate(1, 3), [Node.attr_node(3), Node.attr_node(4), Node.attr_node(5)])\n",
    "t23 = Node.threshold_node(Gate(2, 3), [t33, Node.attr_node(6), t13])\n",
    "\n",
    "print('----------------------------------- Access Tree Structure -----------------------------------')\n",
    "print(t23)\n",
    "print('---------------------------------------------------------------------------------------------')\n",
    "\n",
    "main_keypair = setup(u_size)\n",
    "E = encryption(M, gamma, main_keypair['PK'], u_size)\n",
    "D = key_generation(t23, main_keypair['MK'])\n",
    "assert decrypt_node(E, D, t23), '解密失败！'\n",
    "print('解密成功！解密消息：\\n', M)\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "e5030792b3492f6b12d94f1f48beca3d8e59ec05fd59d0aaaa48e684281ed297"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
